package mat

import (
	"ry/kernel"
)

func (m *floatmatrix) offsetTo(i, j int) int {

	rows, cols := m.backed.Dims()

	return (m.iOffset+i)%rows*m.stride + (m.jOffset+j)%cols
}

func (m *floatmatrix) At(i, j int) float64 { return m.data[m.offsetTo(i, j)] }

func (m *floatmatrix) Set(i, j int, v float64) {
	m.data[m.offsetTo(i, j)] = v
}

func (m *floatmatrix) NumElems() int { return m.rows * m.cols }

func (m *floatmatrix) Dims() (int, int) { return m.rows, m.cols }

func (m *floatmatrix) WrapIndex(i, j int) (int, int) {

	for i < 0 {
		i += m.rows
	}
	for i >= m.rows {
		i -= m.rows
	}
	for j < 0 {
		j += m.cols
	}
	for j >= m.cols {
		j -= m.cols
	}

	return i, j
}

func (m *floatmatrix) Iter() <-chan float64 {

	rows, cols := m.Dims()

	next := make(chan float64)

	go func() {
		for i := 0; i < rows; i++ {
			for j := 0; j < cols; j++ {
				next <- m.At(i, j)
			}
		}
		close(next)
	}()

	return next
}

func (m *floatmatrix) Select(i1, j1, i2, j2 int) Matrix {

	i1, j1 = m.WrapIndex(i1, j1)
	i2, j2 = m.WrapIndex(i2, j2)

	r := *m

	if i2 < i1 { // vertical wrap
		r.rows = m.rows - i1 + i2 + 1
	} else {
		r.rows = i2 - i1 + 1
	}

	if j2 < j1 { // horizontal wrap
		r.cols = m.cols - j1 + j2 + 1
	} else {
		r.cols = j2 - j1 + 1
	}

	r.iOffset = m.iOffset + i1
	r.jOffset = m.jOffset + j1

	return &r
}

func (in *floatmatrix) Convolve(out Matrix, r int, k func(window Matrix, center *float64)) {

	rows, cols := in.Dims()

	kernel.Launch(rows, cols,
		func(i, j int) bool {
			k(in.Select(i-r, j-r, i+r, j+r), &out.data[out.offsetTo(i, j)])
			return true
		})
}
